package com.ruoyi.web.controller.ai;

import com.plexpt.chatgpt.ChatGPTStream;
import com.plexpt.chatgpt.entity.chat.ChatCompletion;
import com.plexpt.chatgpt.entity.chat.Message;
import com.plexpt.chatgpt.util.Proxys;
import com.ruoyi.common.ai.LocalCache;
import com.ruoyi.common.ai.dto.AIChatOpenSseDTO;
import com.ruoyi.common.annotation.Anonymous;
import com.ruoyi.common.base.ai.vo.AIApiKeyCache;
import com.ruoyi.common.constant.SysConfigConstants;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.customize.util.CommonUtil;
import com.ruoyi.system.service.ISysConfigService;
import com.ruoyi.web.ai.AIDynamicKey;
import com.ruoyi.web.ai.listener.CustomizeSseStreamListener;
import com.ruoyi.web.service.IChatMessageService;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Controller;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.DeleteMapping;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.net.Proxy;
import java.util.List;
import java.util.Objects;

/**
 * AI 聊天
 *
 * @author gdk
 * @since 2023-05-15
 */
@Controller
@RequiredArgsConstructor
@RequestMapping("/ai/chat")
public class AIChatController extends BaseController {

    private final IChatMessageService chatMessageService;

    private final ISysConfigService sysConfigService;
    private final AIDynamicKey aiDynamicKey;

    /**
     * 创建sse连接
     *
     * @param dto dto
     * @return {@link SseEmitter}
     * @author gdk
     * @since 2023/5/16 21:09
     */
    @Anonymous
    @CrossOrigin
    @SneakyThrows
    @GetMapping(value = "/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter sseOpen(@Validated AIChatOpenSseDTO dto) {
        // 国内需要代理，国外不需要，设置到本地vpn软件http(s)端口上
        Proxy proxy = Proxys.http("127.0.0.1", 4780);
        SseEmitter sseEmitter = new SseEmitter(-1L);

        AIApiKeyCache key = aiDynamicKey.getKey();
        if (Objects.isNull(key)) {
            sendError("无法连接到模型，请联系管理员", sseEmitter);
            return sseEmitter;
        }

        Integer maxToken = Integer.valueOf(sysConfigService.selectConfigByKey(SysConfigConstants.AI_CHAT_MAX_TOKEN, "2048"));

        // 组装消息，获取上下文
        List<Message> messageList = chatMessageService.initMessageList(dto);
        ChatGPTStream chatGPTStream = ChatGPTStream.builder()
                .timeout(10)
                .proxy(proxy)
                .apiKey(key.getApiKey())
                .apiHost("https://api.openai.com/")
                .build()
                .init();

        CustomizeSseStreamListener listener = new CustomizeSseStreamListener(sseEmitter, aiDynamicKey, key);

        // 回答完成时回调函数
        listener.setOnComplate(msg -> {
            // System.err.println(msg);
            sendEnd(sseEmitter);
        });

        ChatCompletion chatCompletion = ChatCompletion.builder()
                .messages(messageList)
                .stream(true)
                .model(key.getModel())
                .maxTokens(maxToken)
                .build();
        chatGPTStream.streamChatCompletion(chatCompletion, listener);
        return sseEmitter;
    }

    /**
     * 传递报错信息
     *
     * @param msg        信息
     * @param sseEmitter sseEmitter
     * @author gdk
     * @since 2023/5/27 20:33
     */
    public static void sendError(String msg, SseEmitter sseEmitter) {
        try {
            sseEmitter.send(SseEmitter.event().id("notify-error")
                    .data(msg));
            sseEmitter.complete();
        } catch (IOException e) {
            sseEmitter.completeWithError(e);
        }
    }

    /**
     * 传递更新信息
     *
     * @param msg        信息
     * @param sseEmitter sseEmitter
     * @author gdk
     * @since 2023/5/27 20:33
     */
    public static void sendUpdateTrial(String msg, SseEmitter sseEmitter) {
        try {
            sseEmitter.send(SseEmitter.event().id("notify-update-trial")
                    .data(msg));
        } catch (IOException e) {
            sseEmitter.completeWithError(e);
        }
    }

    /**
     * 传递结束信息
     *
     * @param sseEmitter sseEmitter
     * @author gdk
     * @since 2023/5/27 20:33
     */
    public static void sendEnd(SseEmitter sseEmitter) {
        try {
            sseEmitter.send(SseEmitter.event().data("[DONE]"));
        } catch (IOException e) {
            sseEmitter.completeWithError(e);
        } finally {
            sseEmitter.complete();
        }
    }

    /**
     * 关闭sse连接
     *
     * @return {@link SseEmitter}
     * @author gdk
     * @since 2023/5/16 21:10
     */
    @SneakyThrows
    @DeleteMapping("/sse")
    public AjaxResult sseClose() {
        String userId = CommonUtil.getStrUserIdNotNull();

        SseEmitter sse = (SseEmitter) LocalCache.CACHE.get(userId);
        if (sse != null) {
            sse.complete();
            //移除
            LocalCache.CACHE.remove(userId);
        }

        return success();
    }

}
